﻿#pragma once
#include "Index.h"
#include "NumArray.h"

class ICountOfRangeSum
{
public:
	/**
	 * 区间和的个数.
	 * https://leetcode.cn/problems/count-of-range-sum/
	 * 树状数组 线段树 数组 二分查找 分支 有序集合 归并排序
	 *
	 * 给你一个整数数组nums以及两个整数lower和upper。
	 * 求数组中，值位于范围[lower, upper] （包含 lower 和 upper）之内的 区间和的个数 。
	 * 区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。
	 */
	virtual int countRangeSum(std::vector<int>& nums, int lower, int upper) = 0;
};

class CountOfRangeSum
{
public:
	//超时，在最坏情况下nums构建的二叉搜索树会退化为链表
	class MultiSet:public ICountOfRangeSum
	{
	public:
		int countRangeSum(std::vector<int>& nums, int lower, int upper) override
		{
			if (nums.empty())
				return 0;
			const int numsSize = nums.size();
			int ans = 0;

			//在rangeSum中的数额外的加数
			long long offset = 0;
			//遍历时，每次以遍历到的数字为结尾的区间和
			std::multiset<long long> rangeSum;

			for (int i = 0; i < numsSize; ++i)
			{
				//以上一个数字为结尾的所有区间和加上nums[i],变成以nums[i]结尾的区间和
				offset += nums[i];

				//新增以nums[i]自己结尾的区间
				const long long key = nums[i] - offset;
				rangeSum.insert(key);

				//查询当前满足[lower,upper]的区间和数量
				const auto lower_bound = rangeSum.lower_bound(lower - offset);
				const auto upper_bound = rangeSum.upper_bound(upper - offset);
				//累加到结果中
				ans += std::distance(lower_bound, upper_bound);
			}
			return ans;
		}
	};

	//对前缀和数组归并排序 时间:(nlogn) 空间:O(n)
	class MergeSort:public ICountOfRangeSum
	{
	public:
		int countRangeSum(std::vector<int>& nums, int lower, int upper) override
		{
			if (nums.empty() || lower > upper)
				return 0;
			const int numsSize = nums.size();

			//前缀和数组
			std::vector<long> preSums(numsSize, 0);
			preSums[0] = nums.front();
			for (int i = 1; i < numsSize; ++i)
				preSums[i] = preSums[i - 1] + nums[i];

			const int ans = countRangeSumRecursive(preSums, 0, numsSize - 1, lower, upper);
			return ans;
		}
	private:
		int countRangeSumRecursive(std::vector<long>& preSums,int left,int right,int lower,int upper)
		{
			if(left>=right)
				return left==right && preSums[left] >= lower && preSums[left] <= upper ? 1 : 0;

			const int mid = left + ((right - left) >> 1);
			int ans = 0;
			ans += countRangeSumRecursive(preSums, left, mid , lower, upper);
			ans += countRangeSumRecursive(preSums, mid + 1, right, lower, upper);

			//经过子问题处理后，nums[left..mid]有序 nums[mid+1..right]有序
			//满足条件的区间和数量，区间完全在前半部分数组 或 区间完全在后半部分数组 已经在前面的子问题中处理
			//接下来处理 区间的开始在前半部分数组，区间的结尾在后半部分数组的情况

			int finLeft = mid + 1, finRight = mid + 1;//在后半部分数组中满足条件的滑动窗口[finLeft,finRight)，注意是左闭右开
			for(int start=left;start<=mid;++start)
			{
				//先使finLeft满足 nums[finLeft]-nums[start]>=lower
				while (finLeft<=right && preSums[finLeft] - preSums[start] < lower)
					finLeft++;

				//当finRight满足 nums[finRight]-nums[start]>upper 时停止
				finRight = std::max<>(finLeft, finRight);
				while (finRight<=right && preSums[finRight] - preSums[start] <= upper)
					finRight++;

				ans += finRight - finLeft;
			}

			//回溯到父问题前，将两部分有序数组合并
			const int numsRangSize = right - left + 1;
			std::vector<long> merged(numsRangSize);
			int first = left, second = mid + 1, accept = 0;
			while(first<=mid || second<=right)
			{
				if (first > mid)
					merged[accept++] = preSums[second++];
				else if (second > right)
					merged[accept++] = preSums[first++];
				else
				{
					if (preSums[first] < preSums[second])
						merged[accept++] = preSums[first++];
					else
						merged[accept++] = preSums[second++];
				}
			}
			//拷贝排序后的内容
			for (int i = 0; i < numsRangSize; ++i)
				preSums[left + i] = merged[i];

			return ans;
		}
	};
};

#ifdef DEV_TEST
#include <gtest/gtest.h>
TEST(CountRangeSumTest, countRangeSum)
{
	TEST_SOLUTION1(ICountOfRangeSum, CountOfRangeSum::MergeSort, &)
	{
		std::vector<int> nums;
		constexpr int numsSize = 3000;
		constexpr int minVal = -1000;
		constexpr int maxVal = 10000;

		nums = TestingUtil::randRangeInts(numsSize, minVal, maxVal);
		int lower = std::min<>(nums.front(), nums.back());
		int upper = std::max<>(nums.front(), nums.back());

		NumArray::LaziedRangeSum numArray;
		numArray.initNums(nums);

		int expectCount = 0;
		for (int i = 0; i < numsSize; ++i)
			for (int j = i; j < numsSize; ++j)
			{
				const int rangeSum = numArray.sumRange(i, j);
				if (rangeSum >= lower && rangeSum <= upper)
					expectCount++;
			}

		const int resultCount = solution->countRangeSum(nums, lower, upper);

		ASSERT_EQ(resultCount, expectCount)
			<< "nums=" << FormatedUtil::dumpFormated(nums) << std::endl
			<< "lower=" << lower << std::endl
			<< "upper=" << upper << std::endl;
	};
}
#endif